// Created by using on 2/12/22.

#include "dfc_det.h"

// mask color //label
const unsigned char COLOR[19][3] = {
	{20, 20, 20},	// 0 bg
	{20, 220, 20},	// 1 E
	{20, 20, 220},	// 2 L
	{128, 64, 128}, // 3 - road
	{244, 35, 232}, // 4 - sidewalk

};

// 查看细节
void img_show(std::string win_name, cv::Mat &image)
{
	cv::imshow(win_name, image);
	cv::waitKey(0);
	cv::destroyWindow(win_name);
}

/* max func
// float a= 0.009,b=-8.8,c=-7;
// std::cout<<max_float(a, b, c)<<std::endl;
*/
float max_float(float x, float y, float z)
{
	if (x < y)
		x = y;
	if (x < z)
		x = z;
	return x;
}

CrackSeg::CrackSeg(const std::string param_path, const std::string bin_path, std::vector<int> input_shape)
	: param_path(param_path), bin_path(bin_path), input_shape(input_shape)
{
	crack_net.load_param(param_path.c_str());
	crack_net.load_model(bin_path.c_str());
}

CrackSeg::~CrackSeg()
{
	crack_net.clear();
}

cv::Mat CrackSeg::predict_image(cv::Mat &image)
{
	cv::Mat rgbImage;
	cv::resize(image, rgbImage, cv::Size(input_shape[3], input_shape[2]));
	// // cv::cvtColor(image, rgbImage, cv::COLOR_BGR2RGB);
	ncnn::Mat in = ncnn::Mat::from_pixels_resize(image.data, ncnn::Mat::PIXEL_RGB, image.cols, image.rows, input_shape[3], input_shape[2]);
	in.substract_mean_normalize(MEAN_VALS, NORM_VALS);
	ncnn::Extractor ex = crack_net.create_extractor();
	ex.input("in0", in);
	ncnn::Mat out;
	ex.extract("out0", out);

	cv::Mat color(out.h, out.w, CV_8UC3);

	const float *maskmap0 = out.channel(0);
	const float *maskmap1 = out.channel(1);
	const float *maskmap2 = out.channel(2);

	float *out_data = (float *)out.data;

	unsigned char *colordata = color.data;

#pragma omp parallel for
	for (int i = 0; i < out.h; i++)
	{

		for (int j = 0; j < out.w; j++)
		{

			// 类softmax  func
			float max_pr = max_float(maskmap0[i * out.w + j], maskmap1[i * out.w + j], maskmap2[i * out.w + j]);
			// 针对单分类 CV_8UC1
			// color.at<uchar>(i, j) = maskmap2[i * out.w + j] > maskmap0[i * out.w + j] ? 255 : 0;
			int idx = 0;
			if (max_pr == maskmap0[i * out.w + j])
			{
				idx = 0;
			}
			else if (max_pr == maskmap1[i * out.w + j])
			{
				idx = 1;
			}
			else
			{
				idx = 3;
			}

			colordata[(i * out.w + j) * 3 + 0] = COLOR[idx][0];
			colordata[(i * out.w + j) * 3 + 1] = COLOR[idx][1];
			colordata[(i * out.w + j) * 3 + 2] = COLOR[idx][2];
		}
	}

	// 查看mask
	// img_show("color", color);

	// mask与src img 合并
	cv::addWeighted(rgbImage, 0.5, color, 0.5, 0, rgbImage);
	// cv::resize(rgbImage,rgbImage,cv::Size(image.cols, image.rows));
	return rgbImage;
}


void CrackSeg::predict_image(const std::string &src_image_path, const std::string &dst_path)
{
	cv::Mat image = cv::imread(src_image_path);
	cv::Mat segFrame = predict_image(image);
	cv::imwrite(dst_path, segFrame);
}

void CrackSeg::predict_camera()
{
	cv::Mat frame;
	cv::VideoCapture cap;
	int deviceID{0};
	int apiID{cv::CAP_ANY};
	cap.open(deviceID, apiID);
	if (!cap.isOpened())
	{
		std::cout << "Error, cannot open camera!" << std::endl;
		return;
	}
	//--- GRAB AND WRITE LOOP
	std::cout << "Start grabbing" << std::endl
			  << "Press any key to terminate" << std::endl;
	int count{0};
	clock_t start{clock()}, end{0};
	double fps{0};
	for (;;)
	{
		// wait for a new frame from camera and store it into 'frame'
		cap.read(frame);
		// check if we succeeded
		if (frame.empty())
		{
			std::cout << "ERROR! blank frame grabbed" << std::endl;
			break;
		}
		cv::Mat segFrame = predict_image(frame);

		// fps
		++count;
		end = clock();
		fps = count / (float(end - start) / CLOCKS_PER_SEC);
		if (count >= 50)
		{
			count = 0; //防止计数溢出
			start = clock();
		}
		std::cout << "FPS: " << fps << "  Seg Image Number: " << count << "   time consume:" << (float(end - start) / CLOCKS_PER_SEC) << std::endl;
		//设置绘制文本的相关参数
		std::string text{std::to_string(fps)};
		int font_face = cv::FONT_HERSHEY_COMPLEX;
		double font_scale = 1;
		int thickness = 2;
		int baseline;
		cv::Size text_size = cv::getTextSize(text, font_face, font_scale, thickness, &baseline);

		//将文本框居中绘制
		cv::Point origin;
		origin.x = 20;
		origin.y = 20;
		cv::putText(segFrame, text, origin, font_face, font_scale, cv::Scalar(0, 255, 255), thickness, 8, 0);

		// show live and wait for a key with timeout long enough to show images
		imshow("Live", segFrame);
		if (cv::waitKey(5) >= 0)
			break;
	}
	cap.release();
	cv::destroyWindow("Live");
	return;
}